import pandas as pd
import numpy as np
import os
import matplotlib.pyplot as plt
import matplotlib as mpl
from statsmodels.tsa.statespace.sarimax import SARIMAX
from statsmodels.tsa.arima.model import ARIMA
import warnings
from itertools import product
import matplotlib.dates as mdates
from datetime import timedelta

# 设置中文显示
mpl.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
mpl.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

warnings.filterwarnings('ignore')

# 配置路径
EXCEL_PATH = "E:/M/MathModel/23国赛C题/ds/23年6月24到30.xlsx"  # 请确认路径正确
OUTPUT_DIR = "E:/M/MathModel/23国赛C题/ds/7月1日补货量预测"
os.makedirs(OUTPUT_DIR, exist_ok=True)

# 使用您提供的27个单品名称
TARGET_ITEMS = [
    '小米椒(份)', '西兰花', '西峡花菇(1)', '云南生菜(份)', '云南油麦菜(份)', '螺丝椒', '芜湖青椒(1)',
    '紫茄子(2)', '净藕(1)', '长线茄', '竹叶菜', '娃娃菜', '螺丝椒(份)', '双孢菇(盒)',
    '姜蒜小米椒组合装(小份)', '上海青', '奶白菜', '枝江青梗散花', '苋菜', '木耳菜', '洪湖藕带',
    '小皱皮(份)', '红椒(2)', '小青菜(1)', '红薯尖', '菠菜(份)', '菱角'
]


def load_and_prepare_data(file_path):
    """加载并准备数据"""
    print(f"加载数据: {file_path}")
    try:
        df = pd.read_excel(file_path, parse_dates=['销售日期'])

        # 检查必要列是否存在
        required_cols = ['销售日期', '单品名称', '分类名称', '销量(千克)', '损耗率(%)']
        missing_cols = [col for col in required_cols if col not in df.columns]
        if missing_cols:
            print(f"警告: 缺少必要列: {missing_cols}")
            return None, None

        # 转换日期格式
        df['销售日期'] = pd.to_datetime(df['销售日期'])

        print(f"数据加载成功，共 {len(df)} 条记录")
        print(f"日期范围: {df['销售日期'].min()} 至 {df['销售日期'].max()}")
        print(f"包含单品数量: {df['单品名称'].nunique()}")
        print(f"包含分类数量: {df['分类名称'].nunique()}")

        # 计算每个单品的平均损耗率
        item_loss_rates = df.groupby('单品名称')['损耗率(%)'].mean().to_dict()

        return df, item_loss_rates
    except Exception as e:
        print(f"数据加载失败: {str(e)}")
        return None, None


def create_daily_dataset(df, target_items):
    """创建每日单品销量数据集"""
    print("创建每日单品销量数据集...")

    # 只保留目标单品
    df = df[df['单品名称'].isin(target_items)]

    # 按日期和单品分组，计算每日销量
    daily_data = df.groupby(['销售日期', '单品名称', '分类名称'])['销量(千克)'].sum().reset_index()

    # 创建透视表：行为日期，列为单品
    sales_pivot = daily_data.pivot(
        index='销售日期',
        columns='单品名称',
        values='销量(千克)'
    ).fillna(0)

    # 添加缺失的列（确保所有目标单品都在列中）
    for item in target_items:
        if item not in sales_pivot.columns:
            sales_pivot[item] = 0  # 添加缺失的列并填充0

    # 创建分类映射
    category_map = daily_data[['单品名称', '分类名称']].drop_duplicates().set_index('单品名称')['分类名称'].to_dict()

    # 添加日期特征
    sales_pivot['day_of_week'] = sales_pivot.index.dayofweek
    sales_pivot['is_weekend'] = sales_pivot['day_of_week'].apply(lambda x: 1 if x >= 5 else 0)

    print(f"数据集创建成功，包含 {len(sales_pivot)} 天数据和 {len(sales_pivot.columns) - 2} 个单品")
    return sales_pivot, category_map


def forecast_july1_sales(item_name, item_category, historical_series):
    """预测7月1日的销量"""
    try:
        # 确保有足够的数据点
        if len(historical_series) < 3:
            print(f"警告: {item_name} 只有 {len(historical_series)} 个数据点，使用简单平均预测")
            return max(historical_series.mean(), 0.1)  # 确保最小值为0.1

        # 季节性蔬菜使用SARIMA模型
        if item_category in ['花菜类', '辣椒类', '茄类', '食用菌', '水生根茎类']:
            # 使用固定参数的SARIMA模型
            try:
                model = SARIMAX(
                    historical_series,
                    order=(0, 1, 1),
                    seasonal_order=(0, 1, 1, 7),
                    enforce_stationarity=False,
                    enforce_invertibility=False
                )
                results = model.fit(disp=False)

                # 预测7月1日销量
                forecast = results.get_forecast(steps=1)
                return max(forecast.predicted_mean.values[0], 0.1)
            except:
                # SARIMA失败时使用最近3天平均值
                return max(historical_series.iloc[-3:].mean(), 0.1)

        # 花叶类使用ARIMA模型
        else:
            # 使用ARIMA(1,1,1)模型
            try:
                model = ARIMA(historical_series, order=(1, 1, 1))
                results = model.fit()

                # 预测7月1日销量
                forecast = results.forecast(steps=1)
                return max(forecast.values[0], 0.1)
            except:
                # ARIMA失败时使用最近3天平均值
                return max(historical_series.iloc[-3:].mean(), 0.1)

    except Exception as e:
        print(f"预测{item_name}时出错: {str(e)}，使用最近3天平均值")
        return max(historical_series.iloc[-3:].mean(), 0.1)


def calculate_replenishment(predicted_sales, loss_rate):
    """计算考虑损耗率的补货量"""
    if loss_rate >= 100:  # 防止损耗率异常
        loss_rate = 30
    if loss_rate < 0:  # 防止负损耗率
        loss_rate = 0

    # 确保最小销量为0.1kg
    predicted_sales = max(predicted_sales, 0.1)

    # 计算补货量，并确保最小值为0.1kg
    replenishment = max(predicted_sales / (1 - loss_rate / 100), 0.1)
    return replenishment


def plot_sales_trends(sales_data, forecast_results, output_dir):
    """绘制销量趋势图（包含7月1日预测）"""
    print("生成销量趋势图...")
    os.makedirs(output_dir, exist_ok=True)

    # 为每个单品创建图表
    for item in TARGET_ITEMS:
        if item not in sales_data.columns:
            continue

        plt.figure(figsize=(10, 6))

        # 历史销量
        sales_data[item].plot(marker='o', label='历史销量')

        # 7月1日预测销量
        july1_date = sales_data.index[-1] + timedelta(days=1)
        plt.scatter(july1_date, forecast_results[item]['predicted_sales'],
                    color='red', s=100, zorder=5, label='7月1日预测销量')

        # 图表设置
        plt.title(f'{item}销量趋势')
        plt.xlabel('日期')
        plt.ylabel('销量(千克)')
        plt.legend()
        plt.grid(True, linestyle='--', alpha=0.7)

        # 保存图表
        plt.savefig(os.path.join(output_dir, f'{item}_销量趋势.png'), bbox_inches='tight')
        plt.close()

    print(f"销量趋势图已保存至: {output_dir}")


def main():
    """主函数"""
    # 加载数据
    df, item_loss_rates = load_and_prepare_data(EXCEL_PATH)
    if df is None:
        return

    # 创建每日数据集（只针对27个目标单品）
    sales_data, category_map = create_daily_dataset(df, TARGET_ITEMS)

    # 预测每个单品7月1日的销量
    print("\n开始预测各单品7月1日销量...")
    forecast_results = {}

    for item in TARGET_ITEMS:
        # 获取分类
        item_category = category_map.get(item, '未知分类')

        # 获取时间序列数据
        if item in sales_data.columns:
            item_series = sales_data[item]
        else:
            print(f"警告: {item} 无销售数据，使用最近3天平均值")
            # 使用最近3天的平均值（如果可用）
            if len(sales_data) >= 3:
                avg_value = sales_data.iloc[-3:].mean().mean()
            else:
                avg_value = 1.0  # 默认值
            item_series = pd.Series([avg_value] * len(sales_data), index=sales_data.index)

        # 预测7月1日销量
        predicted_sales = forecast_july1_sales(item, item_category, item_series)

        # 获取损耗率（默认10%）
        loss_rate = item_loss_rates.get(item, 10.0)

        # 计算补货量
        replenishment = calculate_replenishment(predicted_sales, loss_rate)

        # 保存结果
        forecast_results[item] = {
            'predicted_sales': predicted_sales,
            'loss_rate': loss_rate,
            'replenishment': replenishment
        }

        print(f"{item}: 预测销量={predicted_sales:.2f}kg, 损耗率={loss_rate:.2f}%, 补货量={replenishment:.2f}kg")

    print(f"预测完成，共预测 {len(forecast_results)} 个单品")

    # 创建结果DataFrame
    results_df = pd.DataFrame({
        '单品名称': TARGET_ITEMS,
        '预测销量(kg)': [forecast_results[item]['predicted_sales'] for item in TARGET_ITEMS],
        '损耗率(%)': [forecast_results[item]['loss_rate'] for item in TARGET_ITEMS],
        '补货量(kg)': [forecast_results[item]['replenishment'] for item in TARGET_ITEMS]
    })

    # 保存结果到Excel
    output_path = os.path.join(OUTPUT_DIR, '7月1日单品补货量预测.xlsx')
    results_df.to_excel(output_path, index=False)
    print(f"结果已保存至: {output_path}")

    # 绘制销量趋势图
    plot_sales_trends(sales_data, forecast_results, OUTPUT_DIR)

    # 打印结果摘要
    print("\n7月1日补货量预测结果:")
    print(results_df[['单品名称', '补货量(kg)']])

    print(f"\n所有结果已保存至: {OUTPUT_DIR}")


if __name__ == "__main__":
    main()